import os

import dotenv
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.config import get_store
from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore

dotenv.load_dotenv()

# 定义长期存储
store = InMemoryStore()
# 添加一些测试数据,users是命名空间, user_123是key,后面的JSON数据是value
store.put(
    ("users",),
    "user_123",
    {
        "name": "GuguLH",
        "age": "26",
    }
)


# 定义工具
@tool(return_direct=True)
def get_user_info(config: RunnableConfig) -> str:
    """查找用户信息"""
    # 获取长期存储,获取到了后,这个存储组件可读也可以写
    _store = get_store()
    # store.put(
    #     ("users",),
    #     "user_456",
    #     {
    #         "name": "Kazamori",
    #         "age": "26",
    #     }
    # )
    # 获取配置中的user_id
    user_id = config["configurable"].get("user_id")
    user_info = _store.get(("users",), user_id)
    return str(user_info.value) if user_info else "Unknown user"


llm = ChatOpenAI(
    model="deepseek-chat",
    base_url=os.getenv("DS_BASE"),
    api_key=os.getenv("DS_API_KEY"),
    temperature=0
).bind_tools([get_user_info])

agent = create_react_agent(
    model=llm,
    tools=[get_user_info],
    store=store,
)

print(agent.invoke(
    {"messages": [{"role": "user", "content": "查找用户信息"}]},
    config={"configurable": {"user_id": "user_123"}},
))
